package com.authsphere.security.authorization;

import com.authsphere.security.token.common.AccessToken;
import com.authsphere.security.token.common.AuthenticationTokenServices;
import com.authsphere.security.token.common.codec.DefaultBase64Codec;
import com.authsphere.security.token.common.codec.TokenCodec;
import jakarta.servlet.http.Cookie;
import jakarta.servlet.http.HttpServletRequest;
import jakarta.servlet.http.HttpServletResponse;
import org.apache.commons.logging.Log;
import org.apache.commons.logging.LogFactory;
import org.springframework.beans.factory.InitializingBean;
import org.springframework.context.MessageSource;
import org.springframework.context.MessageSourceAware;
import org.springframework.context.support.MessageSourceAccessor;
import org.springframework.core.log.LogMessage;
import org.springframework.security.authentication.AccountStatusException;
import org.springframework.security.authentication.AuthenticationDetailsSource;
import org.springframework.security.authentication.AuthenticationManager;
import org.springframework.security.core.Authentication;
import org.springframework.security.core.SpringSecurityMessageSource;
import org.springframework.security.core.userdetails.UsernameNotFoundException;
import org.springframework.security.web.authentication.RememberMeServices;
import org.springframework.security.web.authentication.WebAuthenticationDetailsSource;
import org.springframework.security.web.authentication.logout.LogoutHandler;
import org.springframework.security.web.authentication.rememberme.CookieTheftException;
import org.springframework.security.web.authentication.rememberme.InvalidCookieException;
import org.springframework.security.web.authentication.rememberme.RememberMeAuthenticationException;
import org.springframework.util.Assert;
import org.springframework.util.StringUtils;

import java.io.UnsupportedEncodingException;
import java.net.URLDecoder;
import java.net.URLEncoder;
import java.nio.charset.StandardCharsets;
import java.util.Arrays;
import java.util.Base64;

/**
 * @program: AuthSphere
 * @description: 这里我只需要重写onLoginSuccess方法，因为本项目实际上还是SpringCloud,所以autologin的逻辑是在Gateway中实现的，所以这里就不需要了
 * @author: YuKai Fan
 * @create: 2025/3/14 20:54
 **/
public class PersistentTokenRedisRememberMeServices implements RememberMeServices, InitializingBean, LogoutHandler, MessageSourceAware {
    public static final String SPRING_SECURITY_REMEMBER_ME_COOKIE_KEY = "remember-me";

    public static final String DEFAULT_PARAMETER = "remember-me";

    public static final int DEFAULT_RENEWAL_INTERVAL = 600;

    public static final int TWO_WEEKS_S = 1209600;

    private static final String DELIMITER = ":";

    protected final Log logger = LogFactory.getLog(getClass());

    protected MessageSourceAccessor messages = SpringSecurityMessageSource.getAccessor();
    private AuthenticationDetailsSource<HttpServletRequest, ?> authenticationDetailsSource = new WebAuthenticationDetailsSource();

    private String cookieName = SPRING_SECURITY_REMEMBER_ME_COOKIE_KEY;

    private String cookieDomain;

    private String parameter = DEFAULT_PARAMETER;

    private boolean alwaysRemember;

    private int tokenValiditySeconds = TWO_WEEKS_S;

    private Boolean useSecureCookie = null;

    private Boolean autoRenewal = true;

    private int renewalInterval = DEFAULT_RENEWAL_INTERVAL;

    private final AuthorizationTokenAuthenticationProvider authenticationProvider;
    private final AuthenticationTokenServices authenticationTokenServices;

    private TokenCodec tokenCodec = new DefaultBase64Codec();

    public PersistentTokenRedisRememberMeServices(AuthorizationTokenAuthenticationProvider authenticationProvider) {
        this.authenticationProvider = authenticationProvider;
        this.authenticationTokenServices = authenticationProvider.getAuthenticationTokenServices();
    }

    @Override
    public void afterPropertiesSet() {
    }

    /**
     * Template implementation which locates the Spring Security cookie, decodes it into a
     * delimited array of tokens and submits it to subclasses for processing via the
     * <tt>processAutoLoginCookie</tt> method.
     * <p>
     * The returned username is then used to load the UserDetails object for the user,
     * which in turn is used to create a valid authentication token.
     */
    @Override
    public Authentication autoLogin(HttpServletRequest request, HttpServletResponse response) {
        String rememberMeCookie = extractRememberMeCookie(request);
        if (StringUtils.hasText(rememberMeCookie)) {
            this.logger.debug("Cookie was empty");
            cancelCookie(request, response);
            return null;
        }
        try {
            String[] cookieTokens = tokenCodec.decode(rememberMeCookie);
            return processAutoLoginCookie(cookieTokens, request, response);
        } catch (CookieTheftException ex) {
            cancelCookie(request, response);
            throw ex;
        } catch (UsernameNotFoundException ex) {
            this.logger.debug("Remember-me login was valid but corresponding user not found.", ex);
        } catch (InvalidCookieException ex) {
            this.logger.debug("Invalid remember-me cookie: " + ex.getMessage());
        } catch (AccountStatusException ex) {
            this.logger.debug("Invalid UserDetails: " + ex.getMessage());
        } catch (RememberMeAuthenticationException ex) {
            this.logger.debug(ex.getMessage());
        }
        cancelCookie(request, response);
        return null;
    }

    protected Authentication processAutoLoginCookie(String[] cookieTokens, HttpServletRequest request,
                                                    HttpServletResponse response) {
        if (cookieTokens.length != 2) {
            throw new InvalidCookieException("Cookie token did not contain " + 2 + " tokens, but contained '"
                    + Arrays.asList(cookieTokens) + "'");
        }
        String presentedSeries = cookieTokens[0];
        String presentedToken = cookieTokens[1];
        AccessToken token = this.authenticationTokenServices.getAccessToken(presentedSeries);
        if (token == null) {
            // No series match, so we can't authenticate using this cookie
            throw new RememberMeAuthenticationException("No persistent token found for series id: " + presentedSeries);
        }
        // We have a match for this user/series combination
        if (!presentedToken.equals(token.getValue())) {
            // Token doesn't match series value. Delete all logins for this user and throw
            // an exception to warn them.
            this.authenticationTokenServices.removeAccessToken(token.getSeriesId());
            throw new CookieTheftException(this.messages.getMessage(
                    "PersistentTokenBasedRememberMeServices.cookieStolen",
                    "Invalid remember-me token (Series/token) mismatch. Implies previous cookie theft attack."));
        }
        if (token.isExpired()) {
            throw new RememberMeAuthenticationException("Remember-me login has expired");
        }
        final AuthorizationAuthenticationToken authorizationAuthenticationToken = new AuthorizationAuthenticationToken(token, buildAuthorizationRequest());
        return this.authenticationProvider.authenticate(authorizationAuthenticationToken);
    }

    private AuthorizationRequest buildAuthorizationRequest() {
        return new AuthorizationRequest() {
            @Override
            public boolean needAutoRenewal() {
                return autoRenewal;
            }

            @Override
            public int renewalInterval() {
                return renewalInterval;
            }
        };
    }

    /**
     * Locates the Spring Security remember me cookie in the request and returns its
     * value. The cookie is searched for by name and also by matching the context path to
     * the cookie path.
     *
     * @param request the submitted request which is to be authenticated
     * @return the cookie value (if present), null otherwise.
     */
    protected String extractRememberMeCookie(HttpServletRequest request) {
        Cookie[] cookies = request.getCookies();
        if (cookies == null) {
            return null;
        }
        for (Cookie cookie : cookies) {
            if (this.cookieName.equals(cookie.getName())) {
                return cookie.getValue();
            }
        }
        return null;
    }

    /**
     * Decodes the cookie and splits it into a set of token strings using the ":"
     * delimiter.
     *
     * @param cookieValue the value obtained from the submitted cookie
     * @return the array of tokens.
     * @throws InvalidCookieException if the cookie was not base64 encoded.
     */
    protected String[] decodeCookie(String cookieValue) throws InvalidCookieException {
        for (int j = 0; j < cookieValue.length() % 4; j++) {
            cookieValue = cookieValue + "=";
        }
        String cookieAsPlainText;
        try {
            cookieAsPlainText = new String(Base64.getDecoder().decode(cookieValue.getBytes()));
        } catch (IllegalArgumentException ex) {
            throw new InvalidCookieException("Cookie token was not Base64 encoded; value was '" + cookieValue + "'");
        }
        String[] tokens = StringUtils.delimitedListToStringArray(cookieAsPlainText, DELIMITER);
        for (int i = 0; i < tokens.length; i++) {
            tokens[i] = URLDecoder.decode(tokens[i], StandardCharsets.UTF_8);
        }
        return tokens;
    }

    /**
     * Inverse operation of decodeCookie.
     *
     * @param cookieTokens the tokens to be encoded.
     * @return base64 encoding of the tokens concatenated with the ":" delimiter.
     */
    protected String encodeCookie(String[] cookieTokens) {
        StringBuilder sb = new StringBuilder();
        for (int i = 0; i < cookieTokens.length; i++) {
            try {
                sb.append(URLEncoder.encode(cookieTokens[i], StandardCharsets.UTF_8.toString()));
            } catch (UnsupportedEncodingException ex) {
                this.logger.error(ex.getMessage(), ex);
            }
            if (i < cookieTokens.length - 1) {
                sb.append(DELIMITER);
            }
        }
        String value = sb.toString();
        sb = new StringBuilder(new String(Base64.getEncoder().encode(value.getBytes())));
        while (sb.charAt(sb.length() - 1) == '=') {
            sb.deleteCharAt(sb.length() - 1);
        }
        return sb.toString();
    }

    @Override
    public final void loginFail(HttpServletRequest request, HttpServletResponse response) {
        this.logger.debug("Interactive login attempt was unsuccessful.");
        cancelCookie(request, response);
        onLoginFail(request, response);
    }

    protected void onLoginFail(HttpServletRequest request, HttpServletResponse response) {
    }

    /**
     * {@inheritDoc}
     *
     * <p>
     * Examines the incoming request and checks for the presence of the configured
     * "remember me" parameter. If it's present, or if <tt>alwaysRemember</tt> is set to
     * true, calls <tt>onLoginSucces</tt>.
     * </p>
     */
    @Override
    public final void loginSuccess(HttpServletRequest request, HttpServletResponse response,
                                   Authentication successfulAuthentication) {
        if (!rememberMeRequested(request, this.parameter)) {
            this.logger.debug("Remember-me login not requested.");
            return;
        }
        onLoginSuccess(request, response, successfulAuthentication);
    }

    /**
     * Called from loginSuccess when a remember-me login has been requested. Typically
     * implemented by subclasses to set a remember-me cookie and potentially store a
     * record of it if the implementation requires this.
     */
    protected void onLoginSuccess(HttpServletRequest request, HttpServletResponse response,
                                           Authentication successfulAuthentication) {
        String username = successfulAuthentication.getName();
        this.logger.debug(LogMessage.format("Creating new persistent login for user %s", username));
        try {
            AccessToken accessToken = this.authenticationTokenServices.createAccessToken(successfulAuthentication, false);
            addCookie(accessToken, request, response);
        }
        catch (Exception ex) {
            this.logger.error("Failed to save persistent token ", ex);
        }
    }

    private void addCookie(AccessToken token, HttpServletRequest request, HttpServletResponse response) {
        setCookie(new String[] { token.getSeriesId(), token.getValue() }, token.getExpiresIn() == 0 ? getTokenValiditySeconds() : token.getExpiresIn(), request,
                response);
    }


    /**
     * Allows customization of whether a remember-me login has been requested. The default
     * is to return true if <tt>alwaysRemember</tt> is set or the configured parameter
     * name has been included in the request and is set to the value "true".
     *
     * @param request   the request submitted from an interactive login, which may include
     *                  additional information indicating that a persistent login is desired.
     * @param parameter the configured remember-me parameter name.
     * @return true if the request includes information indicating that a persistent login
     * has been requested.
     */
    protected boolean rememberMeRequested(HttpServletRequest request, String parameter) {
        if (this.alwaysRemember) {
            return true;
        }
        String paramValue = request.getParameter(parameter);
        if (paramValue != null) {
            if (paramValue.equalsIgnoreCase("true") || paramValue.equalsIgnoreCase("on")
                    || paramValue.equalsIgnoreCase("yes") || paramValue.equals("1")) {
                return true;
            }
        }
        this.logger.debug(
                LogMessage.format("Did not send remember-me cookie (principal did not set parameter '%s')", parameter));
        return false;
    }

    /**
     * Sets a "cancel cookie" (with maxAge = 0) on the response to disable persistent
     * logins.
     */
    protected void cancelCookie(HttpServletRequest request, HttpServletResponse response) {
        this.logger.debug("Cancelling cookie");
        Cookie cookie = new Cookie(this.cookieName, null);
        cookie.setMaxAge(0);
        cookie.setPath(getCookiePath(request));
        if (this.cookieDomain != null) {
            cookie.setDomain(this.cookieDomain);
        }
        cookie.setSecure((this.useSecureCookie != null) ? this.useSecureCookie : request.isSecure());
        response.addCookie(cookie);
    }

    /**
     * Sets the cookie on the response.
     * <p>
     * By default a secure cookie will be used if the connection is secure. You can set
     * the {@code useSecureCookie} property to {@code false} to override this. If you set
     * it to {@code true}, the cookie will always be flagged as secure. By default the
     * cookie will be marked as HttpOnly.
     *
     * @param tokens   the tokens which will be encoded to make the cookie value.
     * @param maxAge   the value passed to {@link Cookie#setMaxAge(int)}
     * @param request  the request
     * @param response the response to add the cookie to.
     */
    protected void setCookie(String[] tokens, int maxAge, HttpServletRequest request, HttpServletResponse response) {
        String cookieValue = tokenCodec.encode(tokens);
        Cookie cookie = new Cookie(this.cookieName, cookieValue);
        cookie.setMaxAge(maxAge);
        cookie.setPath(getCookiePath(request));
        if (this.cookieDomain != null) {
            cookie.setDomain(this.cookieDomain);
        }
        if (maxAge < 1) {
            cookie.setMaxAge(-1);
        }
        cookie.setSecure((this.useSecureCookie != null) ? this.useSecureCookie : request.isSecure());
        cookie.setHttpOnly(true);
        response.addCookie(cookie);
    }

    private String getCookiePath(HttpServletRequest request) {
        String contextPath = request.getContextPath();
        return (contextPath.length() > 0) ? contextPath : "/";
    }

    /**
     * Implementation of {@code LogoutHandler}. Default behaviour is to call
     * {@code cancelCookie()}.
     */
    @Override
    public void logout(HttpServletRequest request, HttpServletResponse response, Authentication authentication) {
        this.logger.debug(LogMessage
                .of(() -> "Logout of user " + ((authentication != null) ? authentication.getName() : "Unknown")));
        cancelCookie(request, response);
    }

    public void setCookieName(String cookieName) {
        Assert.hasLength(cookieName, "Cookie name cannot be empty or null");
        this.cookieName = cookieName;
    }

    public void setCookieDomain(String cookieDomain) {
        Assert.hasLength(cookieDomain, "Cookie domain cannot be empty or null");
        this.cookieDomain = cookieDomain;
    }

    protected String getCookieName() {
        return this.cookieName;
    }

    public void setAlwaysRemember(boolean alwaysRemember) {
        this.alwaysRemember = alwaysRemember;
    }

    /**
     * Sets the name of the parameter which should be checked for to see if a remember-me
     * has been requested during a login request. This should be the same name you assign
     * to the checkbox in your login form.
     *
     * @param parameter the HTTP request parameter
     */
    public void setParameter(String parameter) {
        Assert.hasText(parameter, "Parameter name cannot be empty or null");
        this.parameter = parameter;
    }

    public String getParameter() {
        return this.parameter;
    }

    public void setTokenValiditySeconds(int tokenValiditySeconds) {
        this.tokenValiditySeconds = tokenValiditySeconds;
    }

    protected int getTokenValiditySeconds() {
        return this.tokenValiditySeconds;
    }

    /**
     * Whether the cookie should be flagged as secure or not. Secure cookies can only be
     * sent over an HTTPS connection and thus cannot be accidentally submitted over HTTP
     * where they could be intercepted.
     * <p>
     * By default the cookie will be secure if the request is secure. If you only want to
     * use remember-me over HTTPS (recommended) you should set this property to
     * {@code true}.
     *
     * @param useSecureCookie set to {@code true} to always user secure cookies,
     *                        {@code false} to disable their use.
     */
    public void setUseSecureCookie(boolean useSecureCookie) {
        this.useSecureCookie = useSecureCookie;
    }

    protected AuthenticationDetailsSource<HttpServletRequest, ?> getAuthenticationDetailsSource() {
        return this.authenticationDetailsSource;
    }

    public void setAuthenticationDetailsSource(
            AuthenticationDetailsSource<HttpServletRequest, ?> authenticationDetailsSource) {
        Assert.notNull(authenticationDetailsSource, "AuthenticationDetailsSource cannot be null");
        this.authenticationDetailsSource = authenticationDetailsSource;
    }

    public void setAutoRenewal(Boolean autoRenewal) {
        this.autoRenewal = autoRenewal;
    }

    public void setRenewalInterval(int renewalInterval) {
        this.renewalInterval = renewalInterval;
    }

    /**
     * @since 5.5
     */
    @Override
    public void setMessageSource(MessageSource messageSource) {
        Assert.notNull(messageSource, "messageSource cannot be null");
        this.messages = new MessageSourceAccessor(messageSource);
    }

    public void setTokenCodec(TokenCodec tokenCodec) {
        this.tokenCodec = tokenCodec;
    }
}
